# -*- coding: UTF-8 -*-
import torch
import config


class Discriminator(torch.nn.Module):
    def __init__(self, n_node, node_emd_init, manifold, c):
        super().__init__()
        self.n_node = n_node
        self.node_emd_init = node_emd_init
        self.manifold = manifold
        self.c = c

        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

        self.embedding_matrix = torch.nn.Parameter(
            torch.tensor(self.node_emd_init, device=self.device, requires_grad=True))
        self.bias_vector = torch.nn.Parameter(torch.zeros(self.n_node, device=self.device))

        self.d_loss = None
        self.node_embedding = None
        self.node_neighbor_embedding = None
        self.bias = None

        self.score = None
        self.score_f = None
        self.reward = None

        self.optimizer = torch.optim.Adam([self.embedding_matrix, self.bias_vector], lr=config.lr_dis)

        self.hyp_node_embedding = None
        self.hyp_node_neighbor_embedding = None
        self.hyp_bias = None
        self.hyp_score = None
        self.hyp_score_f = None

    def train(self, node_id, node_neighbor_id, label):
        self.node_embedding = self.embedding_matrix[node_id]
        self.node_neighbor_embedding = self.embedding_matrix[node_neighbor_id]
        self.bias = self.bias_vector[node_neighbor_id]
        self.bias[torch.isnan(self.bias)] = 0

        # 将欧式特征向量转换为非欧向量
        self.hyp_node_embedding = torch.cat((torch.zeros(len(node_id), 1).to(self.device), self.node_embedding), dim=1)
        self.hyp_node_embedding = self.manifold.expmap0(self.hyp_node_embedding, self.c)
        self.hyp_node_neighbor_embedding = torch.cat(
            (torch.zeros(len(node_neighbor_id), 1).to(self.device), self.node_neighbor_embedding), dim=1)
        self.hyp_node_neighbor_embedding = self.manifold.expmap0(self.hyp_node_neighbor_embedding, self.c)
        self.hyp_bias = self.manifold.proj_tan0(self.bias.view(1, -1), self.c)

        # 进行非欧计算
        self.hyp_score = self.manifold.mobius_matvec(self.hyp_node_embedding, self.hyp_node_neighbor_embedding, self.c)
        self.hyp_score = self.manifold.proj(self.hyp_score, self.c)
        self.hyp_score = self.manifold.mobius_add(self.hyp_score, self.hyp_bias, c=self.c)
        self.hyp_score = self.manifold.proj(self.hyp_score, self.c)

        # 转换回欧式向量
        self.score = self.manifold.logmap0(self.hyp_score, self.c)
        self.score = torch.sum(self.score, 1)
        self.score = torch.sigmoid(self.score).type(torch.FloatTensor)
        self.score = torch.cuda.FloatTensor(self.score.cuda())

        self.optimizer.zero_grad()
        loss = torch.nn.BCELoss()
        l1_loss = torch.sum(loss(self.score, torch.cuda.FloatTensor(label)))
        l2 = config.lambda_dis * (0.5 * torch.sum(self.node_neighbor_embedding ** 2) +
                                  0.5 * torch.sum(self.node_embedding ** 2) +
                                  0.5 * torch.sum(self.bias ** 2))

        self.d_loss = l1_loss + l2
        self.d_loss.backward()
        self.optimizer.step()

        return self.d_loss

    def forward(self, node_id, node_neighbor_id):
        """ calculate reward, by using the weights that were updated after each training epoch of Discriminator"""
        with torch.no_grad():
            self.node_embedding = self.embedding_matrix[node_id]
            self.node_neighbor_embedding = self.embedding_matrix[node_neighbor_id]
            self.bias = self.bias_vector[node_neighbor_id]
            self.bias[torch.isnan(self.bias)] = 0
            # self.score_f = torch.sum(self.node_embedding * self.node_neighbor_embedding, 1) + self.bias

            # 将欧式向量转换为非欧向量
            self.hyp_node_embedding = torch.cat((torch.zeros(len(node_id), 1).to(self.device), self.node_embedding),
                                                dim=1)
            self.hyp_node_embedding = self.manifold.expmap0(self.hyp_node_embedding, self.c)
            self.hyp_node_neighbor_embedding = torch.cat(
                (torch.zeros(len(node_neighbor_id), 1).to(self.device), self.node_neighbor_embedding), dim=1)
            self.hyp_node_neighbor_embedding = self.manifold.expmap0(self.hyp_node_neighbor_embedding, self.c)
            self.hyp_bias = self.manifold.proj_tan0(self.bias.view(1, -1), self.c)

            self.hyp_score_f = self.manifold.mobius_matvec(self.hyp_node_embedding, self.hyp_node_neighbor_embedding,
                                                           self.c)
            self.hyp_score_f = self.manifold.proj(self.hyp_score_f, self.c)
            self.hyp_score_f = self.manifold.mobius_add(self.hyp_score_f, self.hyp_bias, c=self.c)
            self.hyp_score_f = self.manifold.proj(self.hyp_score_f, self.c)

            self.score_f = self.manifold.logmap0(self.hyp_score_f, self.c)
            self.score_f = torch.sum(self.score_f, 1)
            self.score_f = torch.clamp(self.score_f, -10, 10)
            self.reward = torch.log(1 + torch.exp(self.score_f))
            return self.reward
